#include <iostream>
#include <cstring>
#include <chrono>
#include <algorithm>
#include <cmath>
#include <random>
#include <glm/mat4x4.hpp>
#include <glm/gtx/transform.hpp>
#include "VK_UniformBuffer.h"
#include "VK_Context.h"
#include "VK_Image.h"
#include "VK_Texture.h"
#include "VK_Pipeline.h"
#include "VK_DynamicState.h"
#include "VK_Window.h"
#include "../Node.h"
#include "../ObjectFactory.h"

#define H 4
#define W 4

struct UBO {
    glm::mat4 projection;
    glm::mat4 view;
    glm::mat4 model;
};

std::vector<float> vertices = {
    0, 0, 0.5,
    0, -0.5, -0.5,
    0, 0.5, -0.5
};

struct PushData {
    glm::vec4 pos;
    glm::vec4 color;
};

class pushcontants: public Node
{
public:
    std::vector<PushData> pushData;
public:
    pushcontants()
    {
        nodeName = "pushcontants";
    }
    void initPosAndColor()
    {

        std::default_random_engine rndGen((unsigned)time(nullptr));
        std::uniform_real_distribution<float> rndDist(0.0f, 1.0f);

        for (int i = 0; i < H; i++) {
            for (int j = 0; j < W; j++) {
                pushData.push_back({{0.0f, i * 2, j * 2, 0.0f}, {rndDist(rndGen), rndDist(rndGen), rndDist(rndGen), 0.0f}});
            }
        }

    }
public:
    RunningState run()override
    {
        initPosAndColor();
        auto window = VK_Window::getInstance()->getWindow();
        auto title = getNodeTitle();
        glfwSetWindowTitle(window, title.data());

        context = createVkContext();
        VK_Context::VK_Config vkConfig;
        vkConfig.usingDepthAttachment = true;
        bool ok = context->initVulkanDevice(vkConfig);
        if (!ok)
            return RunningState_Failed;

        auto shaderSet = context->createShaderSet();
        shaderSet->addShader("../shader/pushContants/vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addShader("../shader/pushContants/frag.spv", VK_SHADER_STAGE_FRAGMENT_BIT);

        shaderSet->appendVertexAttributeDescription(0, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT, 0);
        shaderSet->appendVertexInputBindingDescription(3 * sizeof(float), 0, VK_VERTEX_INPUT_RATE_VERTEX);
        VkDescriptorSetLayoutBinding uniformBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(0,
                                                                                                     VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addDescriptorSetLayoutBinding(uniformBinding);

        if (!shaderSet->isValid()) {
            std::cerr << "invalid shaderSet" << std::endl;
            shaderSet->release();
            context->release();
            return RunningState_Failed;
        }

        auto ubo = shaderSet->addUniformBuffer(0, sizeof(UBO));
        ubo->setWriteDataCallback([&](char *&data, uint32_t size)->uint32_t {return updateUniformBufferData(data, size);});

        context->initVulkanContext();
        int32_t width, height;
        glfwGetFramebufferSize(window, &width, &height);

        for (int i = 0; i < H * W; i++) {

            auto pipeline = context->createPipeline(shaderSet);
            pipeline->addPushConstant({VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof (PushData)}, (
                                          char *)(
                                          &pushData[i]));
            pipeline->create();
            auto buffer = context->createVertexBuffer(vertices, 3);
            pipeline->addRenderBuffer(buffer);
        }

        context->createCommandBuffers();
        auto result = context->run();
        setAverageFps(context->getFps());
        context->release();
        return result;
    }

    uint32_t updateUniformBufferData(char *&data, uint32_t size)
    {
        assert(sizeof(UBO) == size);
        static auto startTime = std::chrono::high_resolution_clock::now();
        auto currentTime = std::chrono::high_resolution_clock::now();
        float time = std::chrono::duration<float, std::chrono::seconds::period>
                     (currentTime - startTime).count();
        glm::mat4 model = glm::identity<glm::mat4>();
        model *= glm::translate(glm::mat4(1.0f), glm::vec3(0.0f, -3.0f, -2.5f));
        model *= glm::rotate(glm::mat4(1.0f), glm::radians(30.0f) * time, glm::vec3(1.0f, 0.0f, 0.0f));
        auto view = glm::lookAt(glm::vec3(15.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f,
                                                                                                     0.0f, 1.0f));
        auto proj = glm::perspective(glm::radians(45.0f),
                                     context->getSwapChainExtent().width / (float)context->getSwapChainExtent().height, 0.1f, 100.0f);
        proj[1][1] *= -1;

        UBO ubo;
        ubo.projection = proj;
        ubo.model = model;
        ubo.view = view;
        memcpy(data, &ubo, sizeof(UBO));
        return sizeof(UBO);
    }

private:
    VK_Context *context = nullptr;

};

REGISTER_OBJECT(Node, pushcontants)
